#!/usr/bin/env python3
"""
自动化脚本：使用浏览器自动选择快手素材榜账户，并提取有效的 Cookie
"""

import json
import asyncio
import sys
from playwright.async_api import async_playwright


async def get_material_cookie(cookie_str: str, account_id: int = 83764854) -> str:
    """
    使用 Playwright 打开浏览器，在磁力金牛素材榜中选择账户，并返回已初始化的 Cookie
    
    Args:
        cookie_str: 初始 Cookie 字符串
        account_id: 要选择的账户 ID（默认 83764854 - 晚风漫过窗台）
    
    Returns:
        包含有效会话状态的新 Cookie 字符串
    """
    
    async with async_playwright() as p:
        # 启动浏览器 (headless=False 以便观察)
        browser = await p.chromium.launch(headless=False)
        context = await browser.new_context()
        page = await context.new_page()
        
        # 设置初始 Cookie
        print(f"[*] 设置初始 Cookie...")
        cookie_dict = parse_cookie_string(cookie_str)
        for key, value in cookie_dict.items():
            await context.add_cookies([{
                'name': key,
                'value': value,
                'url': 'https://niu.e.kuaishou.com'
            }])
        
        # 访问素材榜页面
        url = f'https://niu.e.kuaishou.com/material/list?__accountId__={account_id}'
        print(f"[*] 打开页面: {url}")
        await page.goto(url, wait_until='networkidle')
        
        # 等待页面加载
        print(f"[*] 等待页面完全加载...")
        await asyncio.sleep(3)
        
        # 调试：打印页面 HTML 片段
        page_content = await page.content()
        if '晚风漫过窗台' in page_content:
            print(f"[✓] 页面包含账户信息")
        else:
            print(f"[!] 页面未包含账户信息，页面可能需要更多时间加载")
        
        # 尝试多种方式查找和点击账户
        print(f"[*] 尝试查找并点击账户...")
        
        clicked = False
        
        # 方法1：查找包含账户名称的按钮或链接
        try:
            elements = await page.query_selector_all('button, a, div[role="button"], [class*="item"]')
            print(f"[*] 找到 {len(elements)} 个可点击元素")
            
            for element in elements:
                text = await element.text_content()
                if text and ('晚风漫过窗台' in text or '83764854' in text):
                    print(f"[*] 找到账户元素: {text.strip()[:50]}")
                    await element.click()
                    print(f"[✓] 已点击账户")
                    clicked = True
                    await asyncio.sleep(2)
                    break
        except Exception as e:
            print(f"[!] 方法1失败: {e}")
        
        # 方法2：使用 XPath 查找
        if not clicked:
            try:
                print(f"[*] 尝试使用 XPath 查找...")
                element = await page.query_selector('xpath=//*[contains(text(), "晚风漫过窗台")]')
                if element:
                    await element.click()
                    print(f"[✓] 通过 XPath 找到并点击了账户")
                    clicked = True
                    await asyncio.sleep(2)
            except Exception as e:
                print(f"[!] XPath 方法失败: {e}")
        
        # 方法3：尝试点击任何账户项（如果列表中有多个）
        if not clicked:
            try:
                print(f"[*] 尝试点击第一个账户项...")
                account_item = await page.query_selector('[class*="account"], [class*="Account"], li')
                if account_item:
                    await account_item.click()
                    print(f"[✓] 已点击第一个账户项")
                    clicked = True
                    await asyncio.sleep(2)
            except Exception as e:
                print(f"[!] 点击账户项失败: {e}")
        
        if not clicked:
            print(f"[!] 自动查找账户失败")
            print(f"[*] 浏览器已打开，请在浏览器中手动选择账户 '晚风漫过窗台' (ID: 83764854)")
            print(f"[*] 完成后按 Enter 键继续...")
            try:
                input()
            except:
                pass
        
        # 再等等以确保请求完成
        print(f"[*] 等待会话初始化...")
        await asyncio.sleep(3)
        
        # 提取 Cookie
        print(f"[*] 提取会话 Cookie...")
        cookies = await context.cookies()
        new_cookie_str = '; '.join([f"{c['name']}={c['value']}" for c in cookies])
        
        # 关闭浏览器
        await browser.close()
        
        return new_cookie_str


def parse_cookie_string(cookie_str: str) -> dict:
    """解析 Cookie 字符串为字典"""
    cookies = {}
    for item in cookie_str.split(';'):
        if '=' in item:
            key, value = item.strip().split('=', 1)
            cookies[key] = value
    return cookies


async def test_cookie(cookie_str: str) -> bool:
    """测试 Cookie 是否有效"""
    import requests
    import ssl_fix
    
    url = 'https://niu.e.kuaishou.com/rest/esp/promotion/materialCenter/ranking/list'
    payload = {'page': 1, 'size': 20, 'selectDay': 3, 'ordered': 'consume'}
    
    headers = {
        'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36',
        'Content-Type': 'application/json',
        'Cookie': cookie_str,
    }
    
    try:
        resp = requests.post(url, headers=headers, json=payload, timeout=20)
        data = resp.json()
        result = data.get('result')
        
        if result == 1:
            items = data.get('data', {}).get('list', [])
            print(f"✅ Cookie 有效！成功获取 {len(items)} 条素材")
            return True
        else:
            msg = data.get('message', data.get('msg', '未知错误'))
            print(f"❌ Cookie 无效: {msg}")
            return False
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        return False


async def main():
    """主函数"""
    
    # 加载当前配置
    print("[*] 加载 app_config.json...")
    with open('app_config.json', 'r') as f:
        config = json.load(f)
    
    old_cookie = config.get('material_cookies', {}).get('list', [''])[0]
    
    if not old_cookie:
        print("[!] 错误：app_config.json 中没有 material_cookies")
        sys.exit(1)
    
    print(f"[*] 当前 Cookie 长度: {len(old_cookie)} 字符")
    print()
    
    # 使用浏览器获取新 Cookie
    print("[*] 启动浏览器自动化...")
    new_cookie = await get_material_cookie(old_cookie)
    
    print()
    print(f"[*] 新 Cookie 长度: {len(new_cookie)} 字符")
    print()
    
    # 测试新 Cookie
    print("[*] 测试新 Cookie...")
    if await test_cookie(new_cookie):
        # 保存到配置文件
        print()
        print("[*] 保存到 app_config.json...")
        config['material_cookies']['list'][0] = new_cookie
        
        with open('app_config.json', 'w') as f:
            json.dump(config, f, indent=4, ensure_ascii=False)
        
        print("[✓] 成功！Cookie 已更新到 app_config.json")
        print()
        print("[提示] 现在您可以运行爬虫来抓取素材榜数据了")
        return 0
    else:
        print("[!] 新 Cookie 测试失败，未保存")
        return 1


if __name__ == '__main__':
    exit_code = asyncio.run(main())
    sys.exit(exit_code)
